from sklearn import linear_model
import numpy as np

from machine_learning.wm_data import load_data


def check_error(x1, x2, *, threshold=0.2):
    sum_error = 0
    counter = 0
    for i, j in zip(x1.flatten(), x2.flatten()):
        sum_error += abs(i - j)
        counter += 1
    avg = sum_error / counter
    return 0 if avg < threshold else avg


def linear():
    data = load_data()
    x, y = data.X, data.Y
    ps, ns = data.Ps, data.Ns
    # linear_inv
    reg = linear_model.LinearRegression()
    reg.fit(x, y)
    w = np.asarray(reg.coef_).transpose()
    b = np.asarray([reg.intercept_]).transpose()
    x_ = np.hstack((x, np.ones((17, 1))))
    w_ = np.vstack((w, b))
    print(check_error(np.matmul(x_, w_), y))
    # pinv
    w_pinv = np.matmul(np.linalg.pinv(x_), y)
    print(check_error(w_, w_pinv, threshold=0.0001))
    # LDA
    s0 = np.cov(ps.transpose())
    s1 = np.cov(ns.transpose())
    sw = s0 + s1
    sw_pinv = np.linalg.pinv(sw)
    u0 = np.mean(ps, axis=0)
    u1 = np.mean(ns, axis=0)
    w = np.matmul(sw_pinv, u0 - u1)
    print(np.matmul(w, ps.transpose()))
    print(np.matmul(w, ns.transpose()))


if __name__ == '__main__':
    linear()
